package com.mfs.proxy.core.client.processor;

import com.mfs.common.utills.ChannelUtil;
import com.mfs.proxy.core.client.config.ClientConfig;
import com.mfs.proxy.core.server.ProxyServer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.util.Iterator;

/**
 * 代理处理器 提供process的默认实现
 *
 * @author mfs
 * @date 2023/7/9 22:25
 **/
public class DefaultProxyProcessor implements ProxyProcessor {
    private final static Logger logger = LoggerFactory.getLogger(ProxyServer.class);
    private Transfer requestTransfer;
    private Transfer responseTransfer;

    public DefaultProxyProcessor(Transfer requestTransfer, Transfer responseTransfer) {
        if (requestTransfer == null) {
            requestTransfer = new DoNothingTransfer();
        }
        if (responseTransfer == null) {
            responseTransfer = new DoNothingTransfer();
        }
        this.requestTransfer = requestTransfer;
        this.responseTransfer = responseTransfer;
    }

    /**
     * 默认实现，不做任何处理，做默认的转发
     * */
    @Override
    public void process(SocketChannel proxySocketChannel) {
        Selector selector = null;
        try {
            selector = Selector.open();
        } catch (IOException e) {
            logger.error("打开筛选器失败", e);
            ChannelUtil.closeSocketChannel(proxySocketChannel);
            return;
        }

        SocketChannel targetChannel = null;
        try {
            targetChannel = SocketChannel.open(new InetSocketAddress(ClientConfig.getInstance().getTargetIp(),
                    ClientConfig.getInstance().getTargetPort()));
            targetChannel.configureBlocking(false);
        } catch (IOException e) {
            logger.error("连接目标服务失败", e);
            ChannelUtil.closeSocketChannel(targetChannel, proxySocketChannel);
            return;
        }
        SelectionKey proxyKey = null;
        SelectionKey targetKey = null;
        try {
            proxyKey = proxySocketChannel.register(selector, SelectionKey.OP_READ);
            targetKey = targetChannel.register(selector, SelectionKey.OP_READ);
            while (!Thread.currentThread().isInterrupted()) {
                selector.select();
                Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();
                while (iterator.hasNext()) {
                    SelectionKey next = iterator.next();
                    byte[] content = ChannelUtil.read((SocketChannel)next.channel());
                    if (content.length < 1) {
                        ChannelUtil.closeSocketChannel(targetChannel, proxySocketChannel);
                        targetKey.cancel();
                        proxyKey.cancel();
                        return;
                    }
                    if (next == proxyKey) {
                        requestTransfer.translate(ByteBuffer.wrap(content));
                        while (requestTransfer.hasNext()) {
                            targetChannel.write(requestTransfer.next());
                        }
                    } else if (next == targetKey) {
                        responseTransfer.translate(ByteBuffer.wrap(content));
                        while (responseTransfer.hasNext()) {
                            proxySocketChannel.write(responseTransfer.next());
                        }
                    }
                    iterator.remove();
                }
            }
        } catch (Exception e) {
            logger.error("处理代理请求失败", e);
            ChannelUtil.closeSocketChannel(targetChannel, proxySocketChannel);
        }
    }

    class DoNothingTransfer implements Transfer {
        private ByteBuffer byteBuffer = null;
        @Override
        public void translate(ByteBuffer byteBuffer) {
            this.byteBuffer = byteBuffer;
        }

        @Override
        public boolean hasNext() {
            return byteBuffer != null;
        }

        @Override
        public ByteBuffer next() {
            try {
                return byteBuffer;
            } finally {
                byteBuffer = null;
            }
        }
    }
}
